import pandas as pd
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
import re
import logging
from typing import List, Tuple

class TextProcessor:
    """
    文本处理器 - 处理电影标题、标签等文本数据
    """
    
    def __init__(self, max_features: int = 1000):
        self.max_features = max_features
        self.logger = logging.getLogger(__name__)
        self.tfidf_vectorizer = None
        self.tag_vectorizer = None
    
    def preprocess_title(self, title: str) -> str:
        """预处理电影标题"""
        if pd.isna(title):
            return ""
        
        # 移除年份信息
        title = re.sub(r'\(\d{4}\)', '', title)
        
        # 移除特殊字符，保留字母、数字、空格
        title = re.sub(r'[^a-zA-Z0-9\s]', ' ', title)
        
        # 转换为小写并移除多余空格
        title = title.lower().strip()
        title = re.sub(r'\s+', ' ', title)
        
        return title
    
    def extract_title_features(self, movies_df: pd.DataFrame) -> Tuple[pd.DataFrame, TfidfVectorizer]:
        """提取标题TF-IDF特征"""
        if movies_df is None or 'title' not in movies_df.columns:
            return None, None
        
        # 预处理标题
        processed_titles = movies_df['title'].apply(self.preprocess_title)
        
        # 创建TF-IDF特征
        self.tfidf_vectorizer = TfidfVectorizer(
            max_features=self.max_features,
            stop_words='english',
            ngram_range=(1, 2)
        )
        
        title_features = self.tfidf_vectorizer.fit_transform(processed_titles)
        
        # 转换为DataFrame
        feature_names = self.tfidf_vectorizer.get_feature_names_out()
        title_features_df = pd.DataFrame(
            title_features.toarray(),
            columns=[f'title_tfidf_{i}' for i in range(len(feature_names))],
            index=movies_df.index
        )
        
        self.logger.info(f"提取标题特征: {title_features_df.shape}")
        return title_features_df, self.tfidf_vectorizer
    
    def process_tags(self, tags_df: pd.DataFrame, movies_df: pd.DataFrame) -> pd.DataFrame:
        """处理标签数据，创建电影标签特征"""
        if tags_df is None or movies_df is None:
            return None
        
        # 聚合每个电影的标签
        movie_tags = tags_df.groupby('movieId')['tag'].apply(
            lambda x: ' '.join(x.astype(str))
        ).reset_index()
        
        # 合并到电影数据
        movie_tags_features = movies_df.merge(movie_tags, on='movieId', how='left')
        movie_tags_features['tag'] = movie_tags_features['tag'].fillna('')
        
        # 创建标签TF-IDF特征
        self.tag_vectorizer = TfidfVectorizer(
            max_features=500,  # 标签数量通常较少
            stop_words='english'
        )
        
        tag_features = self.tag_vectorizer.fit_transform(movie_tags_features['tag'])
        
        # 转换为DataFrame
        tag_feature_names = self.tag_vectorizer.get_feature_names_out()
        tag_features_df = pd.DataFrame(
            tag_features.toarray(),
            columns=[f'tag_tfidf_{i}' for i in range(len(tag_feature_names))],
            index=movie_tags_features.index
        )
        
        self.logger.info(f"提取标签特征: {tag_features_df.shape}")
        return tag_features_df
    
    def create_text_features(self, movies_df: pd.DataFrame, tags_df: pd.DataFrame = None) -> pd.DataFrame:
        """创建所有文本特征"""
        text_features = {}
        
        # 标题特征
        title_features_df, _ = self.extract_title_features(movies_df)
        if title_features_df is not None:
            text_features['title'] = title_features_df
        
        # 标签特征
        if tags_df is not None:
            tag_features_df = self.process_tags(tags_df, movies_df)
            if tag_features_df is not None:
                text_features['tags'] = tag_features_df
        
        # 合并所有文本特征
        if text_features:
            combined_features = pd.concat(text_features.values(), axis=1)
            self.logger.info(f"合并文本特征: {combined_features.shape}")
            return combined_features
        else:
            return None
        